import tensorflow as tf
from numpy.core.multiarray import dtype


def to_one_hot(batch_labels, class_num):
    batch_size = tf.size(batch_labels)
    batch_labels = tf.expand_dims(batch_labels, 1)
    indices = tf.expand_dims(tf.range(0, batch_size), 1)
    concated = tf.concat([indices, batch_labels], 1)
    onehot_labels = tf.sparse_to_dense(concated, tf.stack([batch_size, class_num]), 1)
    return onehot_labels


with tf.Session() as sess:
    with tf.device("cpu:0"):
        NUM_CLASSES = 10  # 10分类
        labels = [0, 1, 3, 2]  # sample label
        batch_size = tf.size(labels)
        # batch_size = 2
        labels = tf.expand_dims(labels, 1)  # 增加一个维度
        print("labels:{}".format(sess.run(labels)))

        # indices = tf.constant([0, 0, 0, 0], shape=[4, 1])
        # indices = tf.Variable(tf.zeros([4],dtype=tf.int32))

        ##  生成索引
        indices = tf.Variable([0, 1, 2, 3], dtype=tf.int32)
        # indices = tf.expand_dims(tf.range(0, batch_size,1), 1) #生成索引

        # x = tf.Variable([1, 2, 3])
        sess.run(indices.initializer)
        indices = tf.reshape(indices, [4, 1])
        print("indices:{}".format(sess.run(indices)))

        concated = tf.concat([indices, labels], 1)  # 作为拼接
        print("concated:{}".format(sess.run(concated)))
        stack = tf.stack([batch_size, NUM_CLASSES])
        print("stack:{}".format(sess.run(stack)))

        onehot_labels = tf.sparse_to_dense(concated, stack, 1.0, 0.0)
        print("onehot:{}".format(sess.run(onehot_labels)))

        labels = [4, 6, 3, 7, 1, 1, 2]
        one_hot = to_one_hot(labels, 10)
        print("one hot:{}".format(sess.run(one_hot)))
    # sess = tf.Session()
    # labels = [1,3,5,7,9]
    # batch_size = tf.size(labels)
    # labels = tf.expand_dims(labels, 1)
    # indices = tf.expand_dims(tf.range(0, batch_size, 1), 1)
    # concated = tf.concat([indices, labels],1)
    # onehot_labels = tf.sparse_to_dense(concated, tf.stack([batch_size, 10]), 1.0, 0.0)
